#! /usr/bin/bash

# Test on ubuntu1804 & 2004

if [ `id -u` -ne 0 ]; then
    echo "hxg 自用需要Root权限，后果自负"
    exit 1
fi

LOGFILE="/var/log/nvidia_driver_setup.log"
exec > >(tee -a "$LOGFILE") 2>&1

echo "安装Nvidia driver without CUDA and cuDNN."
echo "流程: "
echo "1. 移除所有 Nvidia driver 相关 "
echo "2. 移除 CUDA "
echo "3. Install the specific Nvidia driver selected by the user from a list of available options."
echo "后果自负！"
read -p "骚年你渴望力量吗? [Y/n] " ans
if [[ "$ans" != "Y" && "$ans" != "y" && "$ans" != "" ]]; then
  echo "想想吧！"
  exit 1
fi

/usr/local/cuda/bin/cuda-uninstaller

apt-get remove --purge "*nvidia*" -y

find /usr/lib -iname "*nvidia*" -delete

# Add 'nouveau' to /etc/modules
# echo 'nouveau' | sudo tee -a /etc/modules

rm /etc/X11/xorg.conf

reboot() {
  read -p "Reboot now? [Y/n] " ans
  if [[ "$ans" == "Y" || "$ans" == "y"  || "$ans" == "" ]]; then
    sync; sync; sync; systemctl reboot
  else
    read -n 1 -s -r -p "Press any key to exit..."
    echo
  fi
  exit 0
}

echo "Nvidia drivers and CUDA have been completely removed" 
read -p "来来来，继续? [Y/n] " ans
if [[ "$ans" != "Y" && "$ans" != "y" && "$ans" != "" ]]; then
  echo "咋啦？后悔了"
  echo "See the log file at $LOGFILE"
  reboot
fi

apt-get install nvidia-common

echo | add-apt-repository ppa:graphics-drivers

apt-get update

echo
echo "桥豆麻袋，正在找合适版本的驱动……"
echo
driver_info=$(ubuntu-drivers devices)

devices_info=()
while IFS= read -r line; do
  devices_info+=("$line")
done <<< "$driver_info"

echo "Driver options:"
for ((i=4; i<${#devices_info[@]}; i++)); do
  option="${devices_info[$i]#*: }"
  echo "($((i-3))) $option"
done

while true; do
  read -p "选择一个，开始后续的冒险吧: " choice

  if [[ $choice =~ ^[1-9][0-9]*$ && $choice -le $(( ${#devices_info[@]} - 4 )) ]]; then
    selected_driver=$(echo "${devices_info[$((choice+3))]}" | awk '{for(i=3;i<=NF;i++)printf "%s ", $i; print ""}')
    pkg_name=$(echo "${devices_info[$((choice+3))]}" | awk '{print $3}')
    read -p "Installing '$selected_driver', is that what you want? [Y/n] " ans

    if [[ "$ans" == "Y" || "$ans" == "y" || "$ans" == "" ]]; then
      break
    fi

  else
    echo "Error: Invalid input."
  fi
done

apt-get install $pkg_name -y

apt-get autoremove
apt-get autoclean

echo
echo "'$selected_driver' 安装成功了. CUDA 和 cuDNN 自行处理吧"
echo "See the log file at $LOGFILE"
echo "现在重启，之后请自行使用'nvidia-smi' 验证本次安装"
reboot
